from EmbeddingModel import Embedding
from VectorDatabase import MilvusMuster
from Config import Milvus
# from pdf_analyze import en_extract_text,extract_text_from_pdf

from fastapi import FastAPI, HTTPException, File, UploadFile
from pydantic import BaseModel
import os
from time import sleep
from image_extract import SESSION, extract_images_from_pdf_url

SAVE_DIRECTORY = "cache"

embedding = Embedding()
passed = True
# milvus = MilvusMuster(host = Milvus.HOST, port = Milvus.PORT, alias = Milvus.ALIAS)
while passed:
    try:
        milvus = MilvusMuster(host = Milvus.HOST, port = Milvus.PORT, alias = Milvus.ALIAS)
        passed = False
    except:
        print("尝试连接...")
        sleep(1)

app = FastAPI()

class SearchItem(BaseModel):
    '''查询的数据结构'''
    article_id: list
    content: str
    limit: int = 20
class SearchItemUser(BaseModel):
    '''查询的数据结构'''
    user_id: str
    content: str
    limit: int = 20
        
class InsertItem(BaseModel):
    '''插入的数据结构'''
    article_id: str
    paragraph_id: str
    content: str
    user_id: str
        
        
@app.post("/embedding/distance_search")
async def distance_search(item: SearchItem) -> dict:
    '''根据问题和文章列表，给出相似段落id'''
    q_vector = embedding.get_embedding(item.content)
    result = milvus.get_top_paragraphs(Q_vector = q_vector, articles = item.article_id, limit = item.limit)
    return result

@app.post("/embedding/distance_search_compare")
async def distance_search_compare(item: SearchItem) -> dict:
    '''根据问题和文章列表，给出相似段落id，多篇文章对比专用'''
    q_vector = embedding.get_embedding(item.content)
    result = {"paragraph_ids":[],"scores":[]}
    for article_id in item.article_id:
        temp = milvus.get_top_paragraphs(Q_vector = q_vector, articles = [article_id], limit = 1)
        result["paragraph_ids"] += temp["paragraph_ids"]
        result["scores"] += temp["scores"]
    return result
    
@app.post("/embedding/distance_search_by_user")
async def distance_search_by_user(item: SearchItemUser) -> dict:
    '''根据用户id，给出相似段落id'''
    q_vector = embedding.get_embedding(item.content)
    result = milvus.get_top_paragraphs_by_user(Q_vector = q_vector, user_id = item.user_id, limit = item.limit)
    return result

@app.post("/extract/images")
async def extract_image(file_url: str) -> dict:
    '''Image 提取'''
    global SESSION
    return {"status":"successes","content":await extract_images_from_pdf_url(session = SESSION, pdf_url = file_url)}


@app.post("/embedding/insert_paragraph")
async def insert_paragraph(item: InsertItem) -> InsertItem:
    '''插入数据'''
    paragraph_vector = embedding.get_embedding(item.content)
    print(item.article_id,item.paragraph_id)
    result = await milvus.insert_vector(paragraph_id = item.paragraph_id, article_id = item.article_id, paragraph_vector = paragraph_vector, user_id = item.user_id)
    return item
    
@app.post("/embedding/delete_paragraph")
async def delete_paragraph(paragraph_id: str) -> bool:
    '''插入数据'''
    return milvus.delete_paragraph_by_id(paragraph_id)

@app.post("/embedding/delete_article")
async def delete_article(article_id: str) -> bool:
    '''插入数据'''
    return milvus.delete_article_by_id(article_id)


# @app.post("/en_analyze/")
# async def upload_file(file: UploadFile = File(...)):
    
#     # 确保保存目录存在
#     if not os.path.exists(SAVE_DIRECTORY):
#         os.makedirs(SAVE_DIRECTORY)

#     # 拼接保存文件的完整路径
#     file_path = os.path.join(SAVE_DIRECTORY, file.filename)

#     # 保存文件到指定目录
#     with open(file_path, "wb") as f:
#         f.write(await file.read())
#     try:
#         out = en_extract_text(file_path)
#         if len("".join(out)) < 20:
#             out = extract_text_from_pdf(file_path)
            
#     except Exception as e:
#         print(file_path,str(e))
#         out = extract_text_from_pdf(file_path)
        
        
        
#     return out

# 运行FastAPI应用
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8731)
